from .onnx_passes import *


class PassManager:

    pass_dict = {'Default': [keep_one_domain]}

    def __init__(self):
        pass

    @classmethod
    def set_pass(cls, model_type, *pass_funcs):
        if model_type not in cls.pass_dict:
            cls.pass_dict[model_type] = []
            cls.pass_dict[model_type].extend(cls.pass_dict['Default'])

        for pass_func in pass_funcs:
            # TODO 1. check whether pass_func is valid?
            # TODO 2. support callable func or string of pass name.
            cls.pass_dict[model_type].append(pass_func)

    @classmethod
    def get_pass(cls, model_type=None):
        if not model_type:
            key = 'Default'
        elif model_type not in cls.pass_dict:
            key = 'Default'
        else:
            key = model_type
        return cls.pass_dict[key]


PASS_MANAGER = PassManager
PASS_MANAGER.set_pass('CenterNet', rename_mod_node)
